import numpy as np
import re,math
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

'''
b: 蓝色
g: 绿色
R: 红色
c: 青色
m: 洋红色
y: 黄色
k: 黑色
w: 白色
'''

# ‘DontCare’: 0, ‘cyclist’: 1, ‘tricycle’: 2, ‘sm allMot’:3, ‘bigMot’: 4, ‘pedestrian’: 5, ‘crowds’: 6, ‘unknown’: 7}
def getcolor(var):
    return {
        0: 'c',
        # 0: 'k',
        1: 'b',
        2: 'g',
        3: 'r',
        4: 'k',
        5: 'm',
        6: 'y',
        7: 'w',
    }.get(var, 'error')


def readFile(filepath):
    with open(filepath, "r") as f:
        while (True):
            yield f.readline().strip()

def generator(iters):
    for iter in iters:
        yield iter

filepath_xyz = "./data_3d_pts_lit/training/pts/0a0ca205-5510-4b75-9ec6-e26ced5b0f38_channelVELO_TOP.csv"
# filepath_xyz = 'data_3d_pts_lit/testing/pts/0a0a780a-7ec9-4fb8-91d2-135348649e4b_channelVELO_TOP.csv'
# filepath_xyz = '/home/leo/Downloads/training_data_set/shapeNet/train_data/02691156/000022.pts'
# filepath_xyz = '/home/leo/Downloads/training_data_set/shapeNet/train_data/02773838/000145.pts'
# filepath_xyz = '/home/leo/Downloads/training_data_set/shapeNet/train_data/02954340/007620.pts'
# filepath_xyz  = '/home/leo/Downloads/training_data_set/shapeNet/train_data/02958343/000107.pts'

filepath_cate = "./data_3d_pts_lit/training/category/0a0ca205-5510-4b75-9ec6-e26ced5b0f38_channelVELO_TOP.csv"
# filepath_cate = 'data_3d_pts_lit/testing/category/0a0a780a-7ec9-4fb8-91d2-135348649e4b_channelVELO_TOP.csv'
# filepath_cate = '/home/leo/Downloads/training_data_set/shapeNet/train_label/02691156/000022.seg'
# filepath_cate = '/home/leo/Downloads/training_data_set/shapeNet/train_label/02773838/000145.seg'
# filepath_cate = '/home/leo/Downloads/training_data_set/shapeNet/train_label/02954340/007620.seg'
# filepath_cate  = '/home/leo/Downloads/training_data_set/shapeNet/train_label/02958343/000107.seg'

split_char = ','
# split_char = ' '

CAP = 100000
XYZ = np.zeros((3, CAP))
color = []
time = 0
for i in readFile(filepath_xyz):
    if (time >= CAP):
        break
    try:
        li = [float(j) for j in re.split(split_char, i)]
        XYZ[0][time] = li[0]
        XYZ[1][time] = li[1]
        XYZ[2][time] = li[2]
    except Exception as ess:
        break
    finally:
        time = time + 1
print("time {}".format(time))

time = 0
all_target_num = 0
for z in readFile(filepath_cate):
    if (time >= CAP):
        break
    if z.strip() == "":
        break
    color.append(getcolor(int(z)))
    if int(z) != 0:
        all_target_num +=1
    time = time + 1

    # print (li)
print("time {}".format(time))
print("all_target_num {}".format(all_target_num))

# remove 0
c_g = generator(color)

new_pts = []
new_color = []
index = 0
time = 0
loaded_target_num = 0
for c in c_g:
    # if c != 'w':
    # if c != 'gg':
    if math.sqrt((XYZ[0][index]*XYZ[0][index] + XYZ[1][index]*XYZ[1][index] + XYZ[2][index]*XYZ[2][index])) < 35 :
        new_color.append(c)
        if c != 'c':
            loaded_target_num += 1
        node = np.zeros(3)
        node[0] = XYZ[0][index]
        node[1] = XYZ[1][index]
        node[2] = XYZ[2][index]
        new_pts.append(node)
        time +=1
    index +=1
print("time {}".format(time))
print("new_color {}".format(len(new_color)))
print("loaded all num {}".format(len(new_pts)))
print("index {}".format(index))
print("loaded_target_num {}".format(loaded_target_num))
print("target load rate {}".format(float(loaded_target_num)/float(all_target_num)))
print("target balance rate {}".format(float(loaded_target_num)/float(len(new_pts))))


new_pts = np.transpose(new_pts)

fig = plt.figure(dpi=120)
ax = fig.add_subplot(111, projection='3d')
plt.title('point cloud')
# ax.scatter(XYZ[0], XYZ[1], XYZ[2], c=color, marker='.', s=2, linewidth=0, alpha=1, cmap='spectral')
ax.scatter(new_pts[0], new_pts[1], new_pts[2], c=new_color, marker='.', s=2, linewidth=0, alpha=1, cmap='spectral')

# ax.set_facecolor((0,0,0))
ax.axis('scaled')
# ax.xaxis.set_visible(False)
# ax.yaxis.set_visible(False)
ax.set_xlabel('X Label')
ax.set_ylabel('Y Label')
ax.set_zlabel('Z Label')
plt.show()

